use anyhow::Result;
use axum_app::database::DatabaseManager;
use std::env;

#[tokio::main]
async fn main() -> Result<()> {
    // 设置日志
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::from_default_env()
                .add_directive(tracing::Level::INFO.into()),
        )
        .init();

    // 确保设置了 OPENAI_API_KEY 环境变量
    if env::var("OPENAI_API_KEY").is_err() {
        println!("请设置 OPENAI_API_KEY 环境变量");
        return Ok(());
    }

    println!("🚀 向量搜索演示开始");

    // 创建带向量搜索功能的数据库管理器
    let db = DatabaseManager::new_with_vector_search("vector_demo.db").await?;

    println!("📚 创建示例知识库条目...");

    // 创建一些示例知识库条目
    let sample_items = vec![
        (
            "什么是人工智能".to_string(),
            Some("人工智能（AI）是计算机科学的一个分支，致力于创建能够执行通常需要人类智能的任务的系统。这包括学习、推理、问题解决、感知和语言理解。".to_string()),
            "text".to_string(),
        ),
        (
            "机器学习基础".to_string(),
            Some("机器学习是人工智能的一个子集，它使计算机能够在没有明确编程的情况下学习和改进。它使用算法来分析数据，识别模式，并做出预测或决策。".to_string()),
            "text".to_string(),
        ),
        (
            "深度学习简介".to_string(),
            Some("深度学习是机器学习的一个子集，使用具有多个层的神经网络来模拟人脑的工作方式。它在图像识别、自然语言处理和语音识别等领域特别有效。".to_string()),
            "text".to_string(),
        ),
        (
            "自然语言处理".to_string(),
            Some("自然语言处理（NLP）是人工智能的一个领域，专注于计算机与人类语言之间的交互。它涉及开发能够理解、解释和生成人类语言的算法和模型。".to_string()),
            "text".to_string(),
        ),
        (
            "计算机视觉".to_string(),
            Some("计算机视觉是人工智能的一个领域，致力于使计算机能够从数字图像或视频中获取有意义的信息。它包括图像识别、物体检测和图像分割等任务。".to_string()),
            "text".to_string(),
        ),
    ];

    // 批量创建知识库条目
    let created_ids = db.batch_add_knowledge_items_with_vector(sample_items).await?;
    println!("✅ 成功创建 {} 个知识库条目", created_ids.len());

    // 等待一下，确保向量索引完成
    tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;

    println!("\n🔍 开始向量搜索测试...");

    // 测试向量搜索
    let search_queries = vec![
        "什么是AI？",
        "如何训练模型？",
        "图像识别技术",
        "语言理解算法",
    ];

    for query in search_queries {
        println!("\n查询: \"{}\"", query);
        
        match db.vector_search_knowledge_items(query, Some(3)).await {
            Ok(results) => {
                println!("找到 {} 个相关结果:", results.len());
                for (i, result) in results.iter().enumerate() {
                    println!("  {}. {} (相似度: {:.3})", 
                        i + 1, 
                        result.knowledge_item.title,
                        result.similarity_score
                    );
                }
            }
            Err(e) => {
                println!("搜索失败: {}", e);
            }
        }
    }

    println!("\n📝 测试全文搜索 (FTS)...");

    // 测试全文搜索
    let fts_queries = vec![
        "人工智能",
        "神经网络",
        "算法",
    ];

    for query in fts_queries {
        println!("\nFTS 查询: \"{}\"", query);
        
        match db.search_knowledge_items_fts(query, Some(3)) {
            Ok(items) => {
                println!("找到 {} 个结果:", items.len());
                for (i, item) in items.iter().enumerate() {
                    println!("  {}. {}", i + 1, item.title);
                }
            }
            Err(e) => {
                println!("FTS 搜索失败: {}", e);
            }
        }
    }

    println!("\n🧪 测试更新和删除操作...");

    // 测试更新知识库条目
    if let Some(&first_id) = created_ids.first() {
        println!("更新知识库条目 ID: {}", first_id);
        
        match db.update_knowledge_item_with_vector(
            first_id,
            Some("人工智能详解"),
            Some("人工智能（Artificial Intelligence, AI）是计算机科学的一个重要分支，旨在创建能够模拟人类智能行为的系统。现代AI包括机器学习、深度学习、自然语言处理等多个子领域。"),
            None,
            None,
        ).await {
            Ok(true) => println!("✅ 成功更新知识库条目"),
            Ok(false) => println!("❌ 知识库条目未找到"),
            Err(e) => println!("❌ 更新失败: {}", e),
        }

        // 再次搜索以验证更新
        println!("\n验证更新后的搜索结果:");
        match db.vector_search_knowledge_items("人工智能详细介绍", Some(1)).await {
            Ok(results) => {
                if let Some(result) = results.first() {
                    println!("更新后的标题: {}", result.knowledge_item.title);
                }
            }
            Err(e) => println!("验证搜索失败: {}", e),
        }
    }

    println!("\n🎉 向量搜索演示完成！");
    println!("数据库文件: vector_demo.db");
    println!("你可以使用 SQLite 工具查看数据库内容");

    Ok(())
}